#include "semantic.h"
#include <stdexcept>

using namespace std;

void SymbolTable::enterScope() {
    scopes.push_back(unordered_map<string, Type>());
}

void SymbolTable::exitScope() {
    if (scopes.empty()) throw runtime_error("无法退出空作用域");
    scopes.pop_back();
}

void SymbolTable::addSymbol(const string& id, Type type) {
    if (scopes.empty()) enterScope();
    auto& current = scopes.back();
    if (current.find(id) != current.end()) {
        throw SemanticError("重复定义: " + id);
    }
    current[id] = type;
}

Type SymbolTable::lookup(const string& id) const {
    for (auto it = scopes.rbegin(); it != scopes.rend(); ++it) {
        auto found = it->find(id);
        if (found != it->end()) return found->second;
    }
    throw SemanticError("未声明的标识符: " + id);
}

Type SemanticAnalyzer::checkExpr(Expr* expr) {
    if (auto e = dynamic_cast<EInt*>(expr)) {
        return Type::TInt;
    }
    if (auto e = dynamic_cast<EId*>(expr)) {
        return symbolTable.lookup(e->id);
    }
    if (auto e = dynamic_cast<EUnOp*>(expr)) {
        Type t = checkExpr(e->expr.get());
        if (e->op == UnOp::Neg && t != Type::TInt) {
            throw SemanticError("一元运算类型错误");
        }
        if (e->op == UnOp::Not && t != Type::TBool) {
            throw SemanticError("一元运算类型错误");
        }
        return (e->op == UnOp::Neg) ? Type::TInt : Type::TBool;
    }
    if (auto e = dynamic_cast<EBinOp*>(expr)) {
        Type t1 = checkExpr(e->left.get());
        Type t2 = checkExpr(e->right.get());
        
        switch (e->op) {
            case BinOp::Add: case BinOp::Sub: case BinOp::Mul: case BinOp::Div: case BinOp::Mod:
                if (t1 != Type::TInt || t2 != Type::TInt) 
                    throw SemanticError("算术运算类型错误");
                return Type::TInt;
                
            case BinOp::Eq: case BinOp::Neq: case BinOp::Lt: case BinOp::Le: case BinOp::Gt: case BinOp::Ge:
                if (t1 != Type::TInt || t2 != Type::TInt) 
                    throw SemanticError("比较运算类型错误");
                return Type::TBool;
                
            case BinOp::And: case BinOp::Or:
                if (t1 != Type::TInt && t1 != Type::TBool) 
                    throw SemanticError("逻辑运算左操作数必须是整数或布尔");
                if (t2 != Type::TInt && t2 != Type::TBool)
                    throw SemanticError("逻辑运算右操作数必须是整数或布尔");
                return Type::TBool;

                
            default: throw SemanticError("未知二元操作");
        }
    }
    if (auto e = dynamic_cast<ECall*>(expr)) {
        Type funcType = symbolTable.lookup(e->id);
        // 简化处理：实际应检查参数类型和数量
        return Type::TInt; // 假设所有函数返回int
    }
    throw SemanticError("未知表达式类型");
}

void SemanticAnalyzer::checkStmt(Stmt* stmt) {
    if (auto s = dynamic_cast<SExpr*>(stmt)) {
        checkExpr(s->expr.get());
        return;
    }
    if (auto s = dynamic_cast<SReturn*>(stmt)) {
        if (s->expr) {
            Type t = checkExpr(s->expr.get());
            if (t != currentReturnType) {
                throw SemanticError("返回类型不匹配");
            }
        } else if (currentReturnType != Type::TVoid) {
            throw SemanticError("非void函数必须返回值");
        }
        return;
    }
    if (auto s = dynamic_cast<SDecl*>(stmt)) {
        Type t = checkExpr(s->expr.get());
        if (t != Type::TInt) {
            throw SemanticError("变量声明必须是int类型");
        }
        symbolTable.addSymbol(s->id, Type::TInt);
        return;
    }
    if (auto s = dynamic_cast<SAssign*>(stmt)) {
        Type varType = symbolTable.lookup(s->id);
        Type exprType = checkExpr(s->expr.get());
        if (varType != exprType) {
            throw SemanticError("赋值类型不匹配");
        }
        return;
    }
    if (auto s = dynamic_cast<SIf*>(stmt)) {
        Type condType = checkExpr(s->cond.get());
        if (condType != Type::TBool && condType != Type::TInt) {
            throw SemanticError("if条件必须是布尔类型");
        }
        symbolTable.enterScope();
        checkStmt(s->thenStmt.get());
        symbolTable.exitScope();
        
        if (s->elseStmt) {
            symbolTable.enterScope();
            checkStmt(s->elseStmt.get());
            symbolTable.exitScope();
        }
        return;
    }
    if (auto s = dynamic_cast<SWhile*>(stmt)) {
        Type condType = checkExpr(s->cond.get());
        if (condType != Type::TBool && condType != Type::TInt) {
            throw SemanticError("while条件必须是布尔类型");
        }
        
        bool wasInLoop = inLoop;
        inLoop = true;
        symbolTable.enterScope();
        checkStmt(s->body.get());
        symbolTable.exitScope();
        inLoop = wasInLoop;
        return;
    }
    if (dynamic_cast<SBreak*>(stmt) || dynamic_cast<SContinue*>(stmt)) {
        if (!inLoop) {
            throw SemanticError("break/continue只能在循环中使用");
        }
        return;
    }
    if (auto s = dynamic_cast<SBlock*>(stmt)) {
        symbolTable.enterScope();
        for (auto& stmt : s->stmts) {
            checkStmt(stmt.get());
        }
        symbolTable.exitScope();
        return;
    }
    throw SemanticError("未知语句类型");
}

void SemanticAnalyzer::checkFunc(FuncDef* func) {
    symbolTable.enterScope();
    
    // 添加参数到符号表
    for (auto& param : func->params) {
        symbolTable.addSymbol(param->id, Type::TInt);
    }
    
    // 设置返回类型
    currentReturnType = (func->returnType == FuncType::TInt) ? Type::TInt : Type::TVoid;
    
    // 检查函数体
    try {
        checkStmt(func->body.get());
    } catch (...) {
        symbolTable.exitScope();
        throw;
    }
    
    symbolTable.exitScope();
}

void SemanticAnalyzer::checkCompUnit(CompUnit* cu) {
    // 首先添加所有函数声明
    for (auto& func : cu->funcs) {
        symbolTable.addSymbol(func->name, Type::TFunc);
    }
    
    // 检查main函数是否存在
    try {
        symbolTable.lookup("main");
    } catch (...) {
        throw SemanticError("缺少main函数");
    }
    
    // 检查每个函数
    for (auto& func : cu->funcs) {
        checkFunc(func.get());
    }
}